import torch
import torch.nn as nn
import numpy as np

class MyLoss(nn.Module):
    def __init__(self, num_anchors, num_predictions_per_anchor, device):
        super(MyLoss, self).__init__()
        self.num_anchors = num_anchors
        self.num_predictions = num_predictions_per_anchor
        self.device = device

    def forward(self, predictions, targets):
        """
        Args:
            predictions: 模型预测结果，形状为 (batch_size, num_anchors, num_predictions_per_anchor * 2)
            targets: 真实标签，形状为 (batch_size, num_targets, 1)，其中 num_targets 每个样本可能不同
        """
        batch_size = predictions.shape[0]
        loss = 0

        for i in range(batch_size):  # 遍历batch中的所有样本
            # 获取样本预测和目标
            pred = predictions[i]
            target = targets[i]

            # 将预测结果分割为置信度和角度
            pred_probs = pred[:, 0::2]
            pred_thetas = pred[:, 1::2]

            # 获取每个 anchor 的正负样本标签
            positive_mask, negative_mask = self.get_anchor_mask(target, pred_thetas)

            # 计算正样本的损失（置信度和角度）
            loss_pos_prob = torch.nn.functional.mse_loss(pred_probs[positive_mask], torch.ones_like(pred_probs[positive_mask]))
            loss_pos_theta = torch.nn.functional.mse_loss(pred_thetas[positive_mask], target[:, 0])

            # 计算负样本的损失（仅置信度）
            loss_neg_prob = torch.nn.functional.mse_loss(pred_probs[negative_mask], torch.zeros_like(pred_probs[negative_mask]))

            # 累计损失
            loss += loss_pos_prob + loss_pos_theta + loss_neg_prob

        # 平均每个样本的损失
        loss /= batch_size
        return loss

    def get_anchor_mask(self, targets, predictions):
        """
        在一个样本中，获取每个 anchor 的正负样本标签

        Args:
            targets: 真实标签，形状为 (num_targets, 1)
            predictions: 预测角度，形状为 (num_anchors, num_predictions_per_anchor)

        Returns:
            positive_mask: 正样本 mask，形状为 (num_anchors, num_predictions_per_anchor)
            negative_mask: 负样本 mask，形状为 (num_anchors, num_predictions_per_anchor)
        """
        # 计算每个目标所属的 anchor 索引
        target_anchor_indices = torch.floor(targets[:, 0] * self.num_anchors / (2 * np.pi)).long()

        # 初始化 mask
        positive_mask = torch.zeros((self.num_anchors, self.num_predictions), dtype=torch.bool, device=self.device)
        negative_mask = torch.ones_like(positive_mask)

        # 统计每个 anchor 中的目标数量
        anchor_target_counts = torch.bincount(target_anchor_indices, minlength=self.num_anchors)

        # 遍历每个 anchor
        for anchor_index in range(self.num_anchors):
            num_targets_in_anchor = anchor_target_counts[anchor_index]
            # 标记正样本
            for i in range(min(num_targets_in_anchor, self.num_predictions)):
                positive_mask[anchor_index, i] = True
                negative_mask[anchor_index, i] = False

        return positive_mask, negative_mask